import numpy as np

#### 0.read data
imgs = np.load('small_dataset_imgs_1000.npy')
labels = np.load('small_dataset_labels_1000.npy')
img_quan = np.round(imgs[0] * 7)
#### 1. load weight
f1_filters      = 32
f1_kernel_size  = 3
f1_stride       = 1
f1_padding      = 0
f1_feature_size = 26

f2_filters      = 32
f2_kernel_size  = 3
f2_stride       = 2
f2_padding      = 0
f2_feature_size = 6

fc1_out_ch = 54
fc2_out_ch = 10

cnn1_mapping = np.load('conv1.weight_quantized.npy').reshape(f1_filters, f1_kernel_size ** 2).transpose(1, 0)
cnn2_mapping = np.load('conv2.weight_quantized.npy').reshape(f2_filters, f2_kernel_size ** 2 * f1_filters).transpose(1, 0)
fc1_mapping = np.load('fc1.weight_quantized.npy').transpose([1,0])
fc2_mapping = np.load('fc2.weight_quantized.npy').transpose([1,0])
weight_list = [img_quan.T,cnn1_mapping,cnn2_mapping,fc1_mapping,fc2_mapping]
##################################1. write weight ################################
                		        ######         ######
                		        ######         ######
                		        ######         ######
                		        ######         ######
                		     ############   ############
                		      ##########     ##########
                		       ########       ########
                		        ######         ######
                		         ###            ###
                		          #              #
###################################################################################
'''
dic =  {0:"0",1:"1",2:"2",3:"3",4:"4",5:"5",6:"6",7:"7",-1:"f" ,-2:"e" ,-3:"d" ,-4:"c" ,-5:"b" ,-6:"a" ,-7:"9",-8:"8"}
dic2 = {0:"0",1:"1",2:"2",3:"3",4:"4",5:"5",6:"6",7:"7",-1:"15",-2:"14",-3:"13",-4:"12",-5:"11",-6:"10",-7:"9",-8:"8"}
mapping_weight = np.zeros([576,128],np.int8)
for i in range(0,9):
    for j in range(0,32):
        mapping_weight[i,j] = cnn1_mapping[i,j]
for i in range(0,288):
    for j in range(32,64):
        mapping_weight[i,j] = cnn2_mapping[i,j-32]
for i in range(0,288):
    for j in range(64,64+54):
        mapping_weight[i,j] = fc1_mapping[i,j-64-54]
for i in range(0,54):
    for j in range(118,128):
        mapping_weight[i,j] = fc2_mapping[i,j-118]
# weight file_data for sim
'''

'''
file_data = ""
for i in range(576):
    for j in range(128):
        file_data+= dic2[mapping_weight[i,j]]+","
    file_data+="\n"
with open("cnn_weight_rom.txt","w",encoding="utf-8") as f:
    f.write(file_data)
# weight file_data for fpga
file_data = "memory_initialization_radix=16;\nmemory_initialization_vector=\n"
for i in range(576):
    for j in range(128):
        file_data+= dic[mapping_weight[i,j]]
    file_data+=",\n"
with open("cnn_weight_rom.coe","w",encoding="utf-8") as f:
    f.write(file_data)
'''

# tmp  = fc2_q_mapping
# tmp2 = tmp.T.astype(np.int8).reshape(-1)
# file_data = ""
# for i in range(tmp2.size):
#     file_data+= str(tmp2[i]) + ","
# # with open("f2.txt","w",encoding="utf-8") as f:
# #     f.write(file_data)
#

for i,tmp in enumerate(weight_list):
    tmp3 = bytearray(tmp.T.astype(np.int8).reshape(-1))
    with open("weight_int8/"+str(i)+".raw","wb") as f:
        f.write(tmp3)

print("finish")